import pandas as pd
from service.das.stock_day_das import StockDateDayDas
from service.das.stock_kline_das import KLineData, StockKLineDas


class StockAnalyseService:

    def calculate_and_store_kline(self, stock_code: str):
        """
        计算并存储指定股票的周K、月K、年K数据。
        :param stock_code: 股票代码。
        """
        # 从 DAS 获取所有日K数据
        stockDayDas = StockDateDayDas()
        day_data = stockDayDas.search_by_day_all(stock_code)

        if not day_data:
            AntLogger.info(f"No day K-line data available for stock_code={stock_code}.")
            return

        # 转换为 DataFrame
        data = pd.DataFrame([day.__dict__ for day in day_data])
        data["trade_date"] = pd.to_datetime(data["trade_date"])

        # 定义需要生成的 K 线类型及分组规则
        kline_types = {
            "W": data["trade_date"].dt.to_period("W").apply(lambda r: r.start_time),
            "M": data["trade_date"].dt.to_period("M").apply(lambda r: r.start_time),
            "Y": data["trade_date"].dt.to_period("Y").apply(lambda r: r.start_time),
        }

        # 初始化 KLine DAS
        klineDas = StockKLineDas()

        # 清除该股票的历史 K 线数据
        klineDas.delete_kline_data(stock_code)

        # 逐个生成并存储每种类型的 K 线数据
        for kline_type, group_data in kline_types.items():
            data["group"] = group_data

            # 聚合计算
            grouped = data.groupby(["stock_code", "group"])
            aggregated = grouped.agg(
                open=("open", "first"),
                high=("high", "max"),
                low=("low", "min"),
                close=("close", "last"),
                vol=("vol", "sum"),
                amount=("amount", "sum"),
            ).reset_index()

            # 存储聚合后的数据
            klineDas.store_aggregated_kline(aggregated, kline_type)
